module weno_square_mod
  use recon_type_mod
  use recon_math_mod
  use poly_utils_mod
  use smooth_indicators_mod

  implicit none

contains

  recursive subroutine weno_square_init(this, nd, sw, swx, swy, swz, xc, yc, zc, dx, dy, dz, ic, jc, kc, is, ie, js, je, ks, ke, mask, id)

    class(recon_type), intent(inout) :: this
    integer , intent(in), optional :: nd
    integer , intent(in), optional :: sw
    integer , intent(in), optional :: swx
    integer , intent(in), optional :: swy
    integer , intent(in), optional :: swz
    real(16), intent(in), optional :: xc(:)
    real(16), intent(in), optional :: yc(:)
    real(16), intent(in), optional :: zc(:)
    real(16), intent(in), optional :: dx
    real(16), intent(in), optional :: dy
    real(16), intent(in), optional :: dz
    integer , intent(in), optional :: ic
    integer , intent(in), optional :: jc
    integer , intent(in), optional :: kc
    integer , intent(in), optional :: is
    integer , intent(in), optional :: ie
    integer , intent(in), optional :: js
    integer , intent(in), optional :: je
    integer , intent(in), optional :: ks
    integer , intent(in), optional :: ke
    integer , intent(in), optional :: mask(:,:,:) ! Cell mask
    integer , intent(in), optional :: id

    integer i, j, k, m, n, sub_ie, sub_je, ingb, isub
    integer max_ngb

    call this%clear()
    
    if(nd/=2)then
      stop 'weno_square_init support only nd==2'
    endif
    
    this%sw      = sw
    this%swx     = sw
    this%swy     = sw
    this%swz     = 1
    this%nd      = nd
    this%nc      = sw**nd
    this%max_ngb = 3**nd-1
    
    max_ngb = this%max_ngb
    
    this%dx = dx
    this%dy = dy
    if(present(dz))then
      this%dz = dz
    else
      this%dz = 1
    endif
    
    if (present(id)) this%id = id

    if(present(is))then; this%is = is; else; this%is = 1       ; endif
    if(present(ie))then; this%ie = ie; else; this%ie = this%swx; endif
    if(present(js))then; this%js = js; else; this%js = 1       ; endif
    if(present(je))then; this%je = je; else; this%je = this%swy; endif
    if(present(ks))then; this%ks = ks; else; this%ks = 1       ; endif
    if(present(ke))then; this%ke = ke; else; this%ke = 1       ; endif
    
    this%eps = 1.e-15

    allocate(this%cell_mask(this%is:this%ie,this%js:this%je,1))
    allocate(this%poly_mask(sw,sw,1))

    if (present(mask)) then
      this%cell_mask = mask
    else
      this%cell_mask = 1
    end if

    ! Set polynomial mask from cell mask.
    if (this%cell_mask(this%is,this%js,1) == 0) then
      this%poly_mask = this%cell_mask(this%ie:this%is:-1,this%je:this%js:-1,:)
    else if (this%cell_mask(this%ie,this%js,1) == 0) then
      this%poly_mask = this%cell_mask(:,this%je:this%js:-1,:)
    else if (this%cell_mask(this%is,this%je,1) == 0) then
      this%poly_mask = this%cell_mask(this%ie:this%is:-1,:,:)
    else
      this%poly_mask = this%cell_mask
    end if

    ! Only check trouble cell on stencil.
    if (this%id == 0) then
      allocate(this%cell_ngb(nd,max_ngb,sw,sw,1))
      allocate(this%cell_ngb_size(sw,sw,1))
      k = 1
      do j = 1, sw
        do i = 1, sw
          if (this%cell_mask(i,j,1) == 1) then
            ingb = 0
            do n = max(this%js, j - 1), min(this%je, j + 1)
              do m = max(this%is, i - 1), min(this%ie, i + 1)
                if (this%cell_mask(m,n,1) == 1 .and. m /= i .and. n /= j) then
                  ingb = ingb + 1
                  this%cell_ngb(1,ingb,i,j,k) = m
                  this%cell_ngb(2,ingb,i,j,k) = n
                end if
              end do
            end do
            this%cell_ngb_size(i,j,k) = ingb
          else
            this%cell_ngb_size(i,j,k) = 0
          end if
        end do
      end do
    end if

    allocate(this%xc(this%is:this%ie))
    allocate(this%yc(this%js:this%je))
    if (present(xc) .and. present(yc)) then
      this%xc = xc
      this%yc = yc
    else
      ! Set coordinates of cells on the large stencil with origin at center.
      do i = this%is, this%ie
        this%xc(i) = -int(sw / 2) + i - 1
      end do
      do j = this%js, this%je
        this%yc(j) = -int(sw / 2) + j - 1
      end do
    end if

    ! Initialize sub-stencils.
    if (this%id == 0) then
      this%sub_sw = int((sw + 1) / 2) ! Set sub-stencil size.
      this%sub_nc = this%sub_sw**nd
      this%ns     = this%sub_nc
      this%nterm  = this%ns
      allocate(this%subs(this%ns))
      isub = 1
      do j = this%js, this%js + this%sub_sw - 1
        do i = this%is, this%is + this%sub_sw - 1
          ! init sub procedures
          this%subs(isub)%init                   => weno_square_init
          this%subs(isub)%calc_recon_matrix      => weno_square_calc_recon_matrix     
          this%subs(isub)%calc_ideal_coefs       => weno_square_calc_ideal_coefs      
          this%subs(isub)%reconstruct            => weno_square_reconstruct           
          this%subs(isub)%trouble_cell_indicator => weno_square_trouble_cell_indicator
      
          sub_ie = i + this%sub_sw - 1
          sub_je = j + this%sub_sw - 1
          call this%subs(isub)%init(nd=nd, sw=this%sub_sw, xc=this%xc(i:sub_ie), yc=this%yc(j:sub_je), &
                                    dx=this%dx, dy=this%dy, dz=this%dz, is=i, ie=sub_ie, js=j, je=sub_je, id=isub, &
                                    mask=this%cell_mask(i:sub_ie,j:sub_je,:))
          select case (nd)
          case (1)
            select case (this%sub_sw)
            case (3)
              this%subs(isub)%smooth_indicator => smooth_indicator_3
            end select
          case (2)
            select case (this%sub_sw)
            case (2)
              this%subs(isub)%smooth_indicator => smooth_indicator_2x2
            case (3)
              this%subs(isub)%smooth_indicator => smooth_indicator_3x3
            case (4)
              this%subs(isub)%smooth_indicator => smooth_indicator_4x4
            case (5)
              this%subs(isub)%smooth_indicator => smooth_indicator_5x5
            case default
              stop 'Unsupported WENO smooth indicator'
            end select
          end select
          isub = isub + 1
        end do
      end do
    end if

    this%initialized = .true.

  end subroutine weno_square_init

  subroutine weno_square_calc_recon_matrix(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    ! Local double double arrays for preserving precision.
    real(16), allocatable, dimension(:,:) :: A, iA
    integer , allocatable :: avaCellIdx(:)
    integer i, j, k, ipt, n, iterm, dxn, dyn
    integer iCell, iAvaCell
    
    ierr = 0

    if (allocated(this%iA            )) deallocate(this%iA            )
    if (allocated(this%poly          )) deallocate(this%poly          )
    if (allocated(this%poly_r16      )) deallocate(this%poly_r16      )
    if (allocated(this%recon_mtx     )) deallocate(this%recon_mtx     )
    if (allocated(this%recon_mtx_r16 )) deallocate(this%recon_mtx_r16 )
    if (allocated(this%dpoly         )) deallocate(this%dpoly         )
    if (allocated(this%dpoly_r16     )) deallocate(this%dpoly_r16     )
    if (allocated(this%drecon_mtx    )) deallocate(this%drecon_mtx    )
    if (allocated(this%drecon_mtx_r16)) deallocate(this%drecon_mtx_r16)

    allocate(this%iA            (this%nc ,this%nc));                                        this%iA             = 0
    allocate(this%poly          (this%npt,this%nc));                                        this%poly           = 0
    allocate(this%poly_r16      (this%npt,this%nc));                                        this%poly_r16       = 0
    allocate(this%recon_mtx     (this%npt,this%nc));                                        this%recon_mtx      = 0
    allocate(this%recon_mtx_r16 (this%npt,this%nc));                                        this%recon_mtx_r16  = 0
    allocate(this%dpoly         (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%dpoly          = 0
    allocate(this%dpoly_r16     (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%dpoly_r16      = 0
    allocate(this%drecon_mtx    (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%drecon_mtx     = 0
    allocate(this%drecon_mtx_r16(this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%drecon_mtx_r16 = 0

    allocate( A        (this%nc,this%nc))
    allocate(iA        (this%nc,this%nc)); iA = 0
    allocate(avaCellIdx(this%nc        )); avaCellIdx = 0

    ! Set index maps.
    iAvaCell = 1
    iCell = 1
    do j = this%js, this%je
      do i = this%is, this%ie
        if (this%cell_mask(i,j,1) == 1) then
          avaCellIdx(iCell) = iAvaCell
          iAvaCell = iAvaCell + 1
        end if
        iCell = iCell + 1
      end do
    end do

    ! Set the poly for each evaluation point.
    ! Select monomials according to mask.
    select case (this%nd)
    case (1)
      do ipt = 1, this%npt
        iterm = 1
        do i = 1, this%sw
          if (this%poly_mask(i,1,1) == 1) then
            call calc_monomial(this%x(ipt), i - 1, this%poly_r16(ipt,iterm))
            iterm = iterm + 1
          end if
        end do
      end do
    case (2)
      do ipt = 1, this%npt
        iterm = 1
        do j = 1, this%sw
          do i = 1, this%sw
            if (this%poly_mask(i,j,1) == 1) then
              call calc_monomial(this%x(ipt), this%y(ipt), i - 1, j - 1, this%poly_r16(ipt,iterm))
              iterm = iterm + 1
            end if
          end do
        end do
      end do
    end select
    
    select case (this%nd)
    case (1)
      do ipt = 1, this%npt
        iterm = 1
        do i = 1, this%sw
          do dxn = 0, this%sw - 1
            call calc_deriv_monomial(this%x(ipt), iterm - 1, 1, this%sw, this%dpoly_r16(ipt,iterm,dxn,0,0))
            this%dpoly_r16(ipt,iterm,dxn,0,0) = this%dpoly_r16(ipt,iterm,dxn,0,0) / this%dx**dxn
          end do
          iterm = iterm + 1
        end do
      end do
    case (2)
      do ipt = 1, this%npt
        iterm = 1
        do j = 1, this%sw
          do i = 1, this%sw
            do dyn = 0, this%swy - 1
              do dxn = 0, this%swx - 1
                call calc_deriv_monomial(this%x(ipt), this%y(ipt), i - 1, j - 1, dxn, dyn, this%sw, this%sw, this%dpoly_r16(ipt,iterm,dxn,dyn,0))
                this%dpoly_r16(ipt,iterm,dxn,dyn,0) = this%dpoly_r16(ipt,iterm,dxn,dyn,0) / ( this%dx**dxn * this%dy**dyn )
              end do
            end do
            iterm = iterm + 1
          end do
        end do
      end do
    end select

    ! Calculate inverse of integral coefficient matrix.
    select case (this%nd)
    case (1)
      call calc_tensor_product_poly_integral_matrix(this%sw, this%xc, A, this%cell_mask(:,1,1), this%poly_mask(:,1,1))
    case (2)
      call calc_tensor_product_poly_integral_matrix(this%sw, this%sw, this%xc, this%yc, A, this%cell_mask(:,:,1), this%poly_mask(:,:,1))
    end select
    n = count(this%cell_mask == 1)
    call inverse_matrix(A(1:n,1:n), iA(1:n,1:n), ierr)
    if (ierr /= 0) then
      print*,'Failed to calculate inverse A in weno_square_calc_recon_matrix'
      deallocate(A, iA, avaCellIdx)
      return
    end if

    this%recon_mtx_r16(:,1:n) = matmul(this%poly_r16(:,1:n), iA(1:n,1:n))

    ! Copy double double iA into double iA.
    this%poly = this%poly_r16
    
    do dyn = 0, this%swy - 1
      do dxn = 0, this%swx - 1
        this%drecon_mtx_r16(:,1:n,dxn,dyn,0) = matmul(this%dpoly_r16(:,1:n,dxn,dyn,0), iA(1:n,1:n))
      end do
    end do

    ! Rearrange iA and recon_mtx.
    do n = 1, this%nc
      if (avaCellIdx(n) /= 0) then
        this%iA        (:,n) = iA                (:,avaCellIdx(n))
        this%recon_mtx (:,n) = this%recon_mtx_r16(:,avaCellIdx(n))
        this%drecon_mtx(:,n,:,:,:) = this%drecon_mtx_r16(:,avaCellIdx(n),:,:,:)
      else
        this%iA        (:,n) = 0
        this%recon_mtx (:,n) = 0
        this%drecon_mtx(:,n,:,:,:) = 0
      end if
    end do

    deallocate(A, iA, avaCellIdx)

  end subroutine weno_square_calc_recon_matrix

  subroutine weno_square_calc_ideal_coefs(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    ! Local double double arrays for preserving precision.
    real(16), allocatable, dimension(:,:) :: A, ATA, iATA, gamma
    integer i, j, k, ipt, m, n, di, dj, isub
    
    real(8), parameter :: theta = 3

    ierr = 0

    if (.not. this%initialized) then
      ierr = 1
      return
    end if
    
    do isub = 1, this%ns
      call this%subs(isub)%calc_recon_matrix(ierr)
      if (ierr /= 0) return
    end do
    call this%calc_recon_matrix(ierr)
    if (ierr /= 0) return

    if (allocated(this%gamma)) deallocate(this%gamma)
    allocate(this%gamma(this%ns,this%npt))

    allocate(    A(this%nc,this%ns )); A = 0
    allocate(  ATA(this%ns,this%ns ))
    allocate( iATA(this%ns,this%ns ))
    allocate(gamma(this%ns,this%npt)); gamma = 0
    do ipt = 1, this%npt
      do isub = 1, this%ns
        di = this%subs(isub)%is - this%is
        dj = this%subs(isub)%js - this%js
        do j = 1, this%sub_sw
          do i = 1, this%sub_sw
            if (this%subs(isub)%cell_mask(this%subs(isub)%is+i-1,this%subs(isub)%js+j-1,1) == 1) then
              m = (j + dj - 1) * this%sw     + i + di
              n = (j      - 1) * this%sub_sw + i
              A(m,isub) = this%subs(isub)%recon_mtx_r16(ipt,n)
            end if
          end do
        end do
      end do

      if( all( A(this%nc/2+1,:)==1 ) )then ! Correct for 3rd order center point
        this%gamma(:,ipt) = 1. / this%ns
      else
        ATA = matmul(transpose(A), A)
        call inverse_matrix(ATA, iATA, ierr)
        if (ierr /= 0) return
        gamma(:,ipt) = matmul(matmul(iATA, transpose(A)), this%recon_mtx_r16(ipt,:))
        ! Set near-zero values to zero exactly.
        do isub = 1, this%ns
          if (abs(gamma(isub,ipt)) < 1.0e-15) gamma(isub,ipt) = 0
        end do
  
        if (abs(sum(matmul(A, gamma(:,ipt)) - this%recon_mtx_r16(1,:))) > 1.0e-15) then
          ierr = 3
          this%gamma(:,ipt) = 0
        else
          this%gamma(:,ipt) = gamma(:,ipt)
        end if
      endif

    end do
    deallocate(A, ATA, iATA, gamma)

    if (allocated(this%rp)) deallocate(this%rp)
    allocate(this%rp(this%ns,this%npt))
    if (allocated(this%rn)) deallocate(this%rn)
    allocate(this%rn(this%ns,this%npt))
    if (allocated(this%sigmap)) deallocate(this%sigmap)
    allocate(this%sigmap(this%npt))
    if (allocated(this%sigman)) deallocate(this%sigman)
    allocate(this%sigman(this%npt))

    this%rp = 0.5d0 * (this%gamma + theta * abs(this%gamma))
    this%rn = this%rp - this%gamma

    this%sigmap = sum(this%rp, 1)
    this%sigman = sum(this%rn, 1)

    do ipt = 1, this%npt
      this%rp(:,ipt) = this%rp(:,ipt) / this%sigmap(ipt)
      this%rn(:,ipt) = this%rn(:,ipt) / this%sigman(ipt)
    end do

  end subroutine weno_square_calc_ideal_coefs

  subroutine weno_square_reconstruct(this, fi, fo, TCI, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out) :: fo(:)     ! Reconstructed function values on evaluation points
    logical, intent(in ), optional :: TCI ! Trouble Cell Indicator(TCI), 1 for existing TC(use WENO), 0 for no TC(use poly)
    integer, intent(out), optional :: ierr

    integer isub, ipt
    real(8) a     (this%nterm,this%ns)
    real(8) beta  (this%ns)
    real(8) alpha (this%ns)

    real(8) fs(this%ns,this%npt)
    real(8) alpha_p(this%ns)
    real(8) alpha_n(this%ns)
    
#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    if(.not.TCI)then
      fo = matmul(this%recon_mtx, pack(fi, .true.))
      return
    endif
    
    ! Calculate point values from each available sub-stencils and smoothness
    ! indicators for those sub-stencils.
    do isub = 1, this%ns
      !a(:,isub) = matmul(this%subs(isub)%iA, pack(fi(this%subs(isub)%is:this%subs(isub)%ie,this%subs(isub)%js:this%subs(isub)%je,this%subs(isub)%ks:this%subs(isub)%ke), .true.))
      a(:,isub) = matmul(this%subs(isub)%iA, pack(fi(this%subs(isub)%is:this%subs(isub)%ie,this%subs(isub)%js:this%subs(isub)%je,1), .true.))
      beta(isub) = this%subs(isub)%smooth_indicator(a(:,isub))
      fs(isub,:) = matmul(this%subs(isub)%poly, a(:,isub))
    end do

    alpha = weno_z_alpha(beta, this%ns, this%eps)

    do ipt = 1, this%npt
      ! Handle negative ideal coefficients by splitting method (Shi et al., 2002).
      alpha_p = this%rp(:,ipt) * alpha
      alpha_n = this%rn(:,ipt) * alpha
      fo(ipt) = this%sigmap(ipt) * sum(alpha_p * fs(:,ipt)) / sum(alpha_p) &
              - this%sigman(ipt) * sum(alpha_n * fs(:,ipt)) / sum(alpha_n)
    end do
      
  end subroutine weno_square_reconstruct

  pure function weno_z_alpha(beta, ns, eps) result(res)

    real(8), intent(in) :: beta(ns)
    integer, intent(in) :: ns
    real(8), intent(in) :: eps
    real(8) res(ns)

    real(8) tau
    integer n, i, j

    tau = 0
    n   = 0
    do j = 1, ns - 1
      do i = j + 1, ns
        tau = tau + abs(beta(i) - beta(j))
        n = n + 1
      end do
    end do
    tau = tau / n

    !tau = (abs(beta(6) - beta(4)) + abs(beta(8) - beta(2)) + abs(beta(3) - beta(7)) + abs(beta(1) - beta(9))) / 4. 
    res = 1 + (tau / (beta + eps))**2

  end function weno_z_alpha

  pure logical function weno_square_trouble_cell_indicator(this, fi, dx) result(res)

    class(recon_type), intent(in) :: this
    real(8), intent(in) :: fi(:,:,:)
    real(8), intent(in) :: dx

    integer i, j, k, m, n, ingb
    real(8), parameter :: sqrt2 = sqrt(2.0d0)
    real(8) c

    res = .false.
    k = 1
    do j = this%js, this%je
      do i = this%is, this%ie
        if (this%cell_ngb_size(i,j,k) > 0) then
          if (fi(i,j,k) == 0) then
            c = 2
          else
            c = 0
            do ingb = 1, this%cell_ngb_size(i,j,k)
              m = this%cell_ngb(1,ingb,i,j,k)
              n = this%cell_ngb(2,ingb,i,j,k)
              c = c + abs(fi(m,n,k) - fi(i,j,k))
            end do
            c = c / (this%cell_ngb_size(i,j,k) * sqrt2 * dx * abs(fi(i,j,k)))
          end if
          if (c >= 1) then
            res = .true.
            return
          end if
        end if
      end do
    end do

  end function weno_square_trouble_cell_indicator
  
  subroutine weno_square_reconstruct_deriv(this, fi, fx, fy, fz, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out), optional :: fx(:)     ! Reconstructed x derivative values on evaluation points
    real(8), intent(out), optional :: fy(:)     ! Reconstructed y derivative values on evaluation points
    real(8), intent(out), optional :: fz(:)     ! Reconstructed y derivative values on evaluation points
    integer, intent(out), optional :: ierr

#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    fx = matmul(this%drecon_mtx(:,:,1,0,0), pack(fi, .true.))
    fy = matmul(this%drecon_mtx(:,:,0,1,0), pack(fi, .true.))
    
  end subroutine weno_square_reconstruct_deriv

end module weno_square_mod
